import pygame
import sys
import math

# 初始化pygame
pygame.init()
pygame.font.init()

# 确保中文能正常显示
font = pygame.font.SysFont(['SimHei', 'WenQuanYi Micro Hei', 'Heiti TC'], 36)
small_font = pygame.font.SysFont(['SimHei', 'WenQuanYi Micro Hei', 'Heiti TC'], 24)

# 游戏窗口设置
WIDTH, HEIGHT = 800, 900
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("中国象棋")

# 颜色定义
BLACK = (0, 0, 0)
WHITE = (255, 255, 255)
RED = (255, 0, 0)
BLUE = (0, 0, 200)
BOARD_COLOR = (205, 133, 63)  # 棋盘颜色（棕色）
GRID_COLOR = (0, 0, 0)  # 网格颜色
SELECTED_COLOR = (255, 215, 0)  # 选中颜色（金色）
VALID_MOVE_COLOR = (0, 255, 0, 100)  # 有效移动位置颜色（半透明绿色）

# 棋盘设置
BOARD_SIZE = 9, 10  # 9列10行
GRID_SIZE = 70  # 网格大小
BOARD_MARGIN_X = (WIDTH - (BOARD_SIZE[0] - 1) * GRID_SIZE) // 2  # 棋盘左右边距
BOARD_MARGIN_Y = (HEIGHT - 100 - (BOARD_SIZE[1] - 1) * GRID_SIZE) // 2  # 棋盘上下边距


# 棋子类型
class PieceType:
    ROOK = 1  # 车
    HORSE = 2  # 马
    ELEPHANT = 3  # 象/相
    ADVISOR = 4  # 士/仕
    GENERAL = 5  # 将/帅
    CANNON = 6  # 炮
    PAWN = 7  # 卒/兵


# 棋子类
class Piece:
    def __init__(self, piece_type, is_red, row, col):
        self.type = piece_type
        self.is_red = is_red
        self.row = row
        self.col = col
        self.selected = False

    def get_name(self):
        names = {
            PieceType.ROOK: '车',
            PieceType.HORSE: '马',
            PieceType.ELEPHANT: '象',
            PieceType.ADVISOR: '士',
            PieceType.GENERAL: '将',
            PieceType.CANNON: '炮',
            PieceType.PAWN: '卒'
        }
        if self.is_red:
            if self.type == PieceType.ELEPHANT:
                return '相'
            if self.type == PieceType.ADVISOR:
                return '仕'
            if self.type == PieceType.GENERAL:
                return '帅'
            if self.type == PieceType.PAWN:
                return '兵'
        return names[self.type]

    def can_move(self, board, to_row, to_col):
        # 不能移动到自己的位置
        if self.row == to_row and self.col == to_col:
            return False

        # 不能移动到自己方的棋子上
        target_piece = board.get_piece(to_row, to_col)
        if target_piece is not None and target_piece.is_red == self.is_red:
            return False

        # 根据不同棋子类型检查移动规则
        if self.type == PieceType.ROOK:
            return self._can_move_rook(board, to_row, to_col)
        elif self.type == PieceType.HORSE:
            return self._can_move_horse(board, to_row, to_col)
        elif self.type == PieceType.ELEPHANT:
            return self._can_move_elephant(board, to_row, to_col)
        elif self.type == PieceType.ADVISOR:
            return self._can_move_advisor(board, to_row, to_col)
        elif self.type == PieceType.GENERAL:
            return self._can_move_general(board, to_row, to_col)
        elif self.type == PieceType.CANNON:
            return self._can_move_cannon(board, to_row, to_col)
        elif self.type == PieceType.PAWN:
            return self._can_move_pawn(board, to_row, to_col)
        return False

    def _can_move_rook(self, board, to_row, to_col):
        # 车只能直走
        if self.row != to_row and self.col != to_col:
            return False

        # 检查路径上是否有棋子
        if self.row == to_row:  # 横向移动
            step = 1 if to_col > self.col else -1
            for c in range(self.col + step, to_col, step):
                if board.get_piece(self.row, c) is not None:
                    return False
        else:  # 纵向移动
            step = 1 if to_row > self.row else -1
            for r in range(self.row + step, to_row, step):
                if board.get_piece(r, self.col) is not None:
                    return False

        return True

    def _can_move_horse(self, board, to_row, to_col):
        # 马走日
        row_diff = abs(to_row - self.row)
        col_diff = abs(to_col - self.col)

        if not ((row_diff == 1 and col_diff == 2) or (row_diff == 2 and col_diff == 1)):
            return False

        # 检查蹩马腿
        if row_diff == 2:
            check_row = self.row + (1 if to_row > self.row else -1)
            check_col = self.col
        else:
            check_row = self.row
            check_col = self.col + (1 if to_col > self.col else -1)

        if board.get_piece(check_row, check_col) is not None:
            return False

        return True

    def _can_move_elephant(self, board, to_row, to_col):
        # 象走田
        row_diff = abs(to_row - self.row)
        col_diff = abs(to_col - self.col)

        if row_diff != 2 or col_diff != 2:
            return False

        # 不能过河
        if (self.is_red and to_row > 4) or (not self.is_red and to_row < 5):
            return False

        # 检查塞象眼
        check_row = (self.row + to_row) // 2
        check_col = (self.col + to_col) // 2

        if board.get_piece(check_row, check_col) is not None:
            return False

        return True

    def _can_move_advisor(self, board, to_row, to_col):
        # 士走斜线，且不能出九宫
        row_diff = abs(to_row - self.row)
        col_diff = abs(to_col - self.col)

        if row_diff != 1 or col_diff != 1:
            return False

        # 检查是否在九宫内
        if to_col < 3 or to_col > 5:
            return False

        if self.is_red:
            if to_row < 7 or to_row > 9:
                return False
        else:
            if to_row < 0 or to_row > 2:
                return False

        return True

    def _can_move_general(self, board, to_row, to_col):
        # 将帅只能在九宫内走一步
        row_diff = abs(to_row - self.row)
        col_diff = abs(to_col - self.col)

        if (row_diff + col_diff) != 1:
            return False

        # 检查是否在九宫内
        if to_col < 3 or to_col > 5:
            return False

        if self.is_red:
            if to_row < 7 or to_row > 9:
                return False
        else:
            if to_row < 0 or to_row > 2:
                return False

        # 检查将帅照面
        if to_col == self.col:
            step = 1 if to_row > self.row else -1
            for r in range(self.row + step, to_row, step):
                if board.get_piece(r, to_col) is not None:
                    return False

            # 检查对面将帅是否在同一列且中间无棋子
            enemy_general = board.find_general(not self.is_red)
            if enemy_general is not None and enemy_general.col == to_col:
                min_row = min(to_row, enemy_general.row)
                max_row = max(to_row, enemy_general.row)
                has_obstacle = False
                for r in range(min_row + 1, max_row):
                    if board.get_piece(r, to_col) is not None:
                        has_obstacle = True
                        break
                if not has_obstacle:
                    return False

        return True

    def _can_move_cannon(self, board, to_row, to_col):
        # 炮的移动类似车，但吃子需要翻山
        if self.row != to_row and self.col != to_col:
            return False

        target_piece = board.get_piece(to_row, to_col)
        obstacle_count = 0

        if self.row == to_row:  # 横向移动
            step = 1 if to_col > self.col else -1
            for c in range(self.col + step, to_col, step):
                if board.get_piece(self.row, c) is not None:
                    obstacle_count += 1
        else:  # 纵向移动
            step = 1 if to_row > self.row else -1
            for r in range(self.row + step, to_row, step):
                if board.get_piece(r, self.col) is not None:
                    obstacle_count += 1

        if target_piece is None:
            # 移动，不能有障碍物
            return obstacle_count == 0
        else:
            # 吃子，必须恰好有一个障碍物
            return obstacle_count == 1

    def _can_move_pawn(self, board, to_row, to_col):
        # 兵卒只能前进或横向移动一步
        row_diff = to_row - self.row
        col_diff = abs(to_col - self.col)

        if col_diff > 1 or abs(row_diff) > 1 or (col_diff == 1 and row_diff != 0):
            return False

        # 检查移动方向
        if self.is_red:
            # 红方兵卒向上移动
            if row_diff > 0:
                return False
            # 未过河只能前进
            if self.row > 4 and col_diff != 0:
                return False
        else:
            # 黑方兵卒向下移动
            if row_diff < 0:
                return False
            # 未过河只能前进
            if self.row < 5 and col_diff != 0:
                return False

        return True


# 棋盘类
class Board:
    def __init__(self):
        self.pieces = []
        self.init_pieces()

    def init_pieces(self):
        # 初始化黑方棋子
        self.pieces.append(Piece(PieceType.ROOK, False, 0, 0))
        self.pieces.append(Piece(PieceType.HORSE, False, 0, 1))
        self.pieces.append(Piece(PieceType.ELEPHANT, False, 0, 2))
        self.pieces.append(Piece(PieceType.ADVISOR, False, 0, 3))
        self.pieces.append(Piece(PieceType.GENERAL, False, 0, 4))
        self.pieces.append(Piece(PieceType.ADVISOR, False, 0, 5))
        self.pieces.append(Piece(PieceType.ELEPHANT, False, 0, 6))
        self.pieces.append(Piece(PieceType.HORSE, False, 0, 7))
        self.pieces.append(Piece(PieceType.ROOK, False, 0, 8))
        self.pieces.append(Piece(PieceType.CANNON, False, 2, 1))
        self.pieces.append(Piece(PieceType.CANNON, False, 2, 7))
        self.pieces.append(Piece(PieceType.PAWN, False, 3, 0))
        self.pieces.append(Piece(PieceType.PAWN, False, 3, 2))
        self.pieces.append(Piece(PieceType.PAWN, False, 3, 4))
        self.pieces.append(Piece(PieceType.PAWN, False, 3, 6))
        self.pieces.append(Piece(PieceType.PAWN, False, 3, 8))

        # 初始化红方棋子
        self.pieces.append(Piece(PieceType.ROOK, True, 9, 0))
        self.pieces.append(Piece(PieceType.HORSE, True, 9, 1))
        self.pieces.append(Piece(PieceType.ELEPHANT, True, 9, 2))
        self.pieces.append(Piece(PieceType.ADVISOR, True, 9, 3))
        self.pieces.append(Piece(PieceType.GENERAL, True, 9, 4))
        self.pieces.append(Piece(PieceType.ADVISOR, True, 9, 5))
        self.pieces.append(Piece(PieceType.ELEPHANT, True, 9, 6))
        self.pieces.append(Piece(PieceType.HORSE, True, 9, 7))
        self.pieces.append(Piece(PieceType.ROOK, True, 9, 8))
        self.pieces.append(Piece(PieceType.CANNON, True, 7, 1))
        self.pieces.append(Piece(PieceType.CANNON, True, 7, 7))
        self.pieces.append(Piece(PieceType.PAWN, True, 6, 0))
        self.pieces.append(Piece(PieceType.PAWN, True, 6, 2))
        self.pieces.append(Piece(PieceType.PAWN, True, 6, 4))
        self.pieces.append(Piece(PieceType.PAWN, True, 6, 6))
        self.pieces.append(Piece(PieceType.PAWN, True, 6, 8))

    def get_piece(self, row, col):
        for piece in self.pieces:
            if piece.row == row and piece.col == col:
                return piece
        return None

    def find_general(self, is_red):
        for piece in self.pieces:
            if piece.type == PieceType.GENERAL and piece.is_red == is_red:
                return piece
        return None

    def move_piece(self, from_row, from_col, to_row, to_col):
        piece = self.get_piece(from_row, from_col)
        if piece is None:
            return False

        target_piece = self.get_piece(to_row, to_col)
        if target_piece is not None:
            if target_piece.is_red == piece.is_red:
                return False
            # 移除被吃掉的棋子
            self.pieces.remove(target_piece)

        # 移动棋子
        piece.row = to_row
        piece.col = to_col
        return True

    def is_checkmate(self, is_red):
        # 检查将/帅是否被吃掉
        general = self.find_general(is_red)
        if general is None:
            return True

        # 简单检查：假设这里实现了检查是否被将死的逻辑
        # 实际实现会很复杂，这里简化处理
        return False


# 游戏类
class Game:
    def __init__(self):
        self.board = Board()
        self.current_turn = True  # True表示红方回合，False表示黑方回合
        self.selected_piece = None
        self.game_over = False
        self.winner = None

    def handle_click(self, row, col):
        if self.game_over:
            return

        piece = self.board.get_piece(row, col)

        # 如果点击了自己的棋子，选中它
        if piece is not None and piece.is_red == self.current_turn:
            self.selected_piece = piece
            return

        # 如果已经选中了棋子，尝试移动
        if self.selected_piece is not None:
            if self.selected_piece.can_move(self.board, row, col):
                if self.board.move_piece(self.selected_piece.row, self.selected_piece.col, row, col):
                    # 检查是否将死对方
                    if self.board.is_checkmate(not self.current_turn):
                        self.game_over = True
                        self.winner = self.current_turn
                    else:
                        # 切换回合
                        self.current_turn = not self.current_turn
                self.selected_piece = None
                return

    def get_valid_moves(self):
        if self.selected_piece is None:
            return []

        moves = []
        for row in range(10):
            for col in range(9):
                if self.selected_piece.can_move(self.board, row, col):
                    moves.append((row, col))
        return moves

    def reset(self):
        self.board = Board()
        self.current_turn = True
        self.selected_piece = None
        self.game_over = False
        self.winner = None


# 绘制棋盘
def draw_board(screen):
    # 绘制背景
    screen.fill(BOARD_COLOR)

    # 绘制边框
    pygame.draw.rect(screen, GRID_COLOR,
                     (BOARD_MARGIN_X - 10, BOARD_MARGIN_Y - 10,
                      (BOARD_SIZE[0] - 1) * GRID_SIZE + 20, (BOARD_SIZE[1] - 1) * GRID_SIZE + 20),
                     5)

    # 绘制网格线
    for i in range(BOARD_SIZE[1]):
        # 横线
        pygame.draw.line(screen, GRID_COLOR,
                         (BOARD_MARGIN_X, BOARD_MARGIN_Y + i * GRID_SIZE),
                         (BOARD_MARGIN_X + (BOARD_SIZE[0] - 1) * GRID_SIZE, BOARD_MARGIN_Y + i * GRID_SIZE),
                         2)

    for i in range(BOARD_SIZE[0]):
        # 竖线
        if i == 0 or i == BOARD_SIZE[0] - 1:
            pygame.draw.line(screen, GRID_COLOR,
                             (BOARD_MARGIN_X + i * GRID_SIZE, BOARD_MARGIN_Y),
                             (BOARD_MARGIN_X + i * GRID_SIZE, BOARD_MARGIN_Y + (BOARD_SIZE[1] - 1) * GRID_SIZE),
                             2)
        else:
            # 河界上方的竖线
            pygame.draw.line(screen, GRID_COLOR,
                             (BOARD_MARGIN_X + i * GRID_SIZE, BOARD_MARGIN_Y),
                             (BOARD_MARGIN_X + i * GRID_SIZE, BOARD_MARGIN_Y + 4 * GRID_SIZE),
                             2)
            # 河界下方的竖线
            pygame.draw.line(screen, GRID_COLOR,
                             (BOARD_MARGIN_X + i * GRID_SIZE, BOARD_MARGIN_Y + 5 * GRID_SIZE),
                             (BOARD_MARGIN_X + i * GRID_SIZE, BOARD_MARGIN_Y + (BOARD_SIZE[1] - 1) * GRID_SIZE),
                             2)

    # 绘制河界
    river_text = "楚河               汉界"
    river_font = font
    text_surface = river_font.render(river_text, True, GRID_COLOR)
    text_rect = text_surface.get_rect(center=(WIDTH // 2, BOARD_MARGIN_Y + 4.5 * GRID_SIZE))
    screen.blit(text_surface, text_rect)

    # 绘制九宫
    # 上方九宫
    pygame.draw.line(screen, GRID_COLOR,
                     (BOARD_MARGIN_X + 3 * GRID_SIZE, BOARD_MARGIN_Y),
                     (BOARD_MARGIN_X + 5 * GRID_SIZE, BOARD_MARGIN_Y + 2 * GRID_SIZE),
                     2)
    pygame.draw.line(screen, GRID_COLOR,
                     (BOARD_MARGIN_X + 5 * GRID_SIZE, BOARD_MARGIN_Y),
                     (BOARD_MARGIN_X + 3 * GRID_SIZE, BOARD_MARGIN_Y + 2 * GRID_SIZE),
                     2)
    # 下方九宫
    pygame.draw.line(screen, GRID_COLOR,
                     (BOARD_MARGIN_X + 3 * GRID_SIZE, BOARD_MARGIN_Y + 7 * GRID_SIZE),
                     (BOARD_MARGIN_X + 5 * GRID_SIZE, BOARD_MARGIN_Y + 9 * GRID_SIZE),
                     2)
    pygame.draw.line(screen, GRID_COLOR,
                     (BOARD_MARGIN_X + 5 * GRID_SIZE, BOARD_MARGIN_Y + 7 * GRID_SIZE),
                     (BOARD_MARGIN_X + 3 * GRID_SIZE, BOARD_MARGIN_Y + 9 * GRID_SIZE),
                     2)

    # 绘制兵卒炮位置标记
    # 黑方兵卒
    for col in [0, 2, 4, 6, 8]:
        # 上方兵卒标记
        draw_pawn_mark(screen, 3, col)
        # 下方兵卒标记
        draw_pawn_mark(screen, 6, col)

    # 炮位置标记
    draw_cannon_mark(screen, 2, 1)
    draw_cannon_mark(screen, 2, 7)
    draw_cannon_mark(screen, 7, 1)
    draw_cannon_mark(screen, 7, 7)


# 绘制兵卒标记
def draw_pawn_mark(screen, row, col):
    x = BOARD_MARGIN_X + col * GRID_SIZE
    y = BOARD_MARGIN_Y + row * GRID_SIZE

    # 右上角
    pygame.draw.line(screen, GRID_COLOR, (x + 10, y), (x + 20, y), 2)
    pygame.draw.line(screen, GRID_COLOR, (x + 20, y), (x + 20, y + 10), 2)

    # 右下角
    pygame.draw.line(screen, GRID_COLOR, (x + 20, y + GRID_SIZE), (x + 20, y + GRID_SIZE - 10), 2)
    pygame.draw.line(screen, GRID_COLOR, (x + 20, y + GRID_SIZE - 10), (x + 10, y + GRID_SIZE), 2)

    # 左下角
    pygame.draw.line(screen, GRID_COLOR, (x, y + GRID_SIZE - 10), (x, y + GRID_SIZE), 2)
    pygame.draw.line(screen, GRID_COLOR, (x, y + GRID_SIZE - 10), (x + 10, y + GRID_SIZE - 10), 2)

    # 左上角
    pygame.draw.line(screen, GRID_COLOR, (x, y), (x, y + 10), 2)
    pygame.draw.line(screen, GRID_COLOR, (x, y), (x + 10, y), 2)


# 绘制炮位置标记
def draw_cannon_mark(screen, row, col):
    x = BOARD_MARGIN_X + col * GRID_SIZE
    y = BOARD_MARGIN_Y + row * GRID_SIZE

    # 右上角
    pygame.draw.line(screen, GRID_COLOR, (x + 5, y), (x + 15, y), 2)
    pygame.draw.line(screen, GRID_COLOR, (x + 15, y), (x + 15, y + 5), 2)

    # 右下角
    pygame.draw.line(screen, GRID_COLOR, (x + 15, y + GRID_SIZE), (x + 15, y + GRID_SIZE - 5), 2)
    pygame.draw.line(screen, GRID_COLOR, (x + 15, y + GRID_SIZE - 5), (x + 5, y + GRID_SIZE), 2)

    # 左下角
    pygame.draw.line(screen, GRID_COLOR, (x, y + GRID_SIZE - 5), (x, y + GRID_SIZE), 2)
    pygame.draw.line(screen, GRID_COLOR, (x, y + GRID_SIZE - 5), (x + 5, y + GRID_SIZE - 5), 2)

    # 左上角
    pygame.draw.line(screen, GRID_COLOR, (x, y), (x, y + 5), 2)
    pygame.draw.line(screen, GRID_COLOR, (x, y), (x + 5, y), 2)


# 绘制棋子
def draw_pieces(screen, game):
    for piece in game.board.pieces:
        x = BOARD_MARGIN_X + piece.col * GRID_SIZE
        y = BOARD_MARGIN_Y + piece.row * GRID_SIZE

        # 绘制棋子背景
        color = RED if piece.is_red else BLUE
        pygame.draw.circle(screen, color, (x, y), 30)
        pygame.draw.circle(screen, WHITE, (x, y), 28)

        # 绘制棋子文字
        text_color = RED if piece.is_red else BLACK
        text_surface = font.render(piece.get_name(), True, text_color)
        text_rect = text_surface.get_rect(center=(x, y))
        screen.blit(text_surface, text_rect)

        # 如果是选中的棋子，绘制高亮
        if piece == game.selected_piece:
            pygame.draw.circle(screen, SELECTED_COLOR, (x, y), 30, 3)

    # 绘制有效移动位置
    if game.selected_piece is not None:
        valid_moves = game.get_valid_moves()
        for row, col in valid_moves:
            x = BOARD_MARGIN_X + col * GRID_SIZE
            y = BOARD_MARGIN_Y + row * GRID_SIZE
            s = pygame.Surface((GRID_SIZE, GRID_SIZE), pygame.SRCALPHA)
            s.fill(VALID_MOVE_COLOR)
            screen.blit(s, (x - GRID_SIZE // 2, y - GRID_SIZE // 2))


# 绘制游戏状态
def draw_game_status(screen, game):
    status_text = f"当前回合: {'红方' if game.current_turn else '黑方'}"
    text_surface = small_font.render(status_text, True, BLACK)
    screen.blit(text_surface, (WIDTH // 2 - text_surface.get_width() // 2, HEIGHT - 60))

    if game.game_over:
        winner_text = f"游戏结束! {'红方' if game.winner else '黑方'}获胜!"
        text_surface = font.render(winner_text, True, BLACK)
        screen.blit(text_surface, (WIDTH // 2 - text_surface.get_width() // 2, HEIGHT - 100))

        restart_text = "按R键重新开始"
        text_surface = small_font.render(restart_text, True, BLACK)
        screen.blit(text_surface, (WIDTH // 2 - text_surface.get_width() // 2, HEIGHT - 30))


# 主函数
def main():
    game = Game()

    running = True
    while running:
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                running = False
            elif event.type == pygame.MOUSEBUTTONDOWN and event.button == 1:
                # 计算点击的棋盘位置
                x, y = event.pos
                col = (x - BOARD_MARGIN_X) // GRID_SIZE
                row = (y - BOARD_MARGIN_Y) // GRID_SIZE

                if 0 <= row < BOARD_SIZE[1] and 0 <= col < BOARD_SIZE[0]:
                    game.handle_click(row, col)
            elif event.type == pygame.KEYDOWN:
                if event.key == pygame.K_r:
                    game.reset()

        # 绘制游戏
        draw_board(screen)
        draw_pieces(screen, game)
        draw_game_status(screen, game)

        # 更新显示
        pygame.display.flip()

    pygame.quit()
    sys.exit()


if __name__ == "__main__":
    main()